import threading
import queue
import numpy as np
import zarr
import os
from typing import Optional, Dict, Any
import dataclasses



    
    
    
# class RobotDataLogger:
#     def __init__(self, path: str = "robot_data.zarr", overwrite: bool = True):
#         if overwrite and os.path.exists(path):
#             import shutil
#             shutil.rmtree(path)

#         self.zarr_root = zarr.open(path, mode='w')
#         self.counter = 0
#         self.lock = threading.Lock()

#         self.q = queue.Queue()
#         self.worker = threading.Thread(target=self._worker_thread, daemon=True)
#         self.worker.start()

#     def _next_sample_id(self) -> str:
#         with self.lock:
#             self.counter += 1
#             return f"sample_{self.counter:06d}"

#     def _worker_thread(self):
#         while True:
#             task = self.q.get()
#             if task is None:
#                 break
#             sample_id, data = task
#             g = self.zarr_root.require_group(sample_id)
#             for key, value in data.items():
#                 if isinstance(value, np.ndarray):
#                     g.array(name=key, data=value, compressor='zstd', overwrite=True)
#                 else:
#                     g.attrs[key] = value

#     def log_sample(
#         self,
#         rgb: Optional[np.ndarray] = None,
#         depth: Optional[np.ndarray] = None,
#         pose: Optional[np.ndarray] = None,
#         imu: Optional[np.ndarray] = None,
#         timestamp: Optional[float] = None,
#         meta: Optional[Dict[str, Any]] = None
#     ):
#         sample_id = self._next_sample_id()
#         data = {}
#         if rgb is not None:
#             data["rgb"] = rgb
#         if depth is not None:
#             data["depth"] = depth
#         if pose is not None:
#             data["pose"] = pose
#         if imu is not None:
#             data["imu"] = imu
#         data["timestamp"] = timestamp if timestamp else get_timestamp()
#         if meta:
#             data.update(meta)

#         self.q.put((sample_id, data))

#     def close(self):
#         self.q.put(None)
#         self.worker.join()


 
 
 
class TrajectoryDataLogger:
    
    def __init__(self, path='demo_data.zarr', start_id=0, overwrite=False):
        self.root = zarr.open(path, mode='w') if overwrite else zarr.open(path, mode='a')
        self.q = queue.Queue()
        self.worker = threading.Thread(target=self._worker, daemon=True)
        self.worker.start()
        self.counter = start_id
        self.lock = threading.Lock()
        self.current_traj = []

    def start_trajectory(self, meta: Optional[dict] = None):
        self.current_traj = []
        self.current_meta = meta or {}

    def record_frame(self, rgb: np.ndarray, force: np.ndarray, timestamp: Optional[float] = None):
        self.current_traj.append({
            'rgb': rgb,
            'force': force,
            'timestamp': timestamp if timestamp else time.time()
        })

    def end_trajectory(self):
        with self.lock:
            self.counter += 1
            traj_id = f"trajectory_{self.counter:06d}"
        self.q.put((traj_id, self.current_traj, self.current_meta))
        self.current_traj = []

    def _worker(self):
        while True:
            task = self.q.get()
            if task is None:
                break
            traj_id, traj_data, meta = task
            g = self.root.require_group(traj_id)

            # Stack and save trajectory
            rgb_seq = np.stack([f['rgb'] for f in traj_data])
            force_seq = np.stack([f['force'] for f in traj_data])
            timestamps = np.array([f['timestamp'] for f in traj_data])

            g.create_dataset('rgb', data=rgb_seq, compressor='zstd', chunks=True, overwrite=True)
            g.create_dataset('force', data=force_seq, compressor='zstd', chunks=True, overwrite=True)
            g.create_dataset('timestamp', data=timestamps, compressor='zstd', chunks=True, overwrite=True)
            g.attrs.update(meta)

    def close(self):
        self.q.put(None)
        self.worker.join()
